import faiss
import sqlite3

def init_faiss(faiss_path, dimension):
    # 初始化 FAISS 索引
    index = faiss.IndexFlatL2(dimension)    # 使用 L2 距离
    index = faiss.IndexIDMap(index)         # 包装成支持 ID 的索引

    # 保存到本地
    faiss.write_index(index, faiss_path)   # 将索引保存到本地 index.faiss 文件中

    return index


def get_cursor(dataset_path):
    # 创建一个sqlite连接
    # 连接到sqlite数据库，如果文件不存在则自动创建
    conn = sqlite3.connect(dataset_path)

    # 创建cursor对象
    cursor = conn.cursor()

    # 创建用户表，如果表不存在
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS origin (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            path TEXT NOT NULL,
            content TEXT NOT NULL,
            tokens INTEGER,
            time REAL
        )
    ''')

    conn.commit()

    return conn, cursor


def get_faiss(faiss_path):
    # 从本地加载索引
    index = faiss.read_index(faiss_path)

    return index


def save_faiss(index, faiss_path):
    faiss.write_index(index, faiss_path) 



# 查询sqlite中的数据
def seach_database():
    conn, cursor = get_cursor("/workspace/database/dataset.db")

    # cursor.execute("delete from origin")
    cursor.execute("select * from origin")

    rows = cursor.fetchall()
    for row in rows:
        print(row)


# 查询faiss中有哪些id
def search_faiss(faiss_path):
    index = get_faiss(faiss_path)

    num_vectors = index.ntotal  # 获取索引中的向量总数
    for i in range(num_vectors):
        print("vector id:", i)


# seach_database()
# search_faiss('/workspace/database/index.faiss')